Add DatabricksOAuthTokenSource by emmyzhou-db · Pull Request #439 · databricks/databricks-sdk-java

emmyzhou-db · 2025-04-30T15:56:46Z

What changes are proposed in this pull request?

Implemented DatabricksOAuthTokenSource class that handles the OAuth token exchange flow
This class manages the OAuth token exchange flow using ID tokens to obtain access tokens
The token exchange mechanism is essential for implementing OIDC-based authentication in Databricks

How is this tested?

The implementation is tested through unit tests that:

Mock the ID Token Source to simulate the token exchange flow
Mock the HTTP client to verify token exchange requests and responses
Test the complete token exchange flow including audience selection and error handling
Validate proper parameter handling and error responses
All tests are automated and part of the unit test suite. No manual testing was required.

NO_CHANGELOG=true

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java

parthban-db · 2025-04-30T21:16:12Z

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java

+
+    OAuthResponse response;
+    try {
+      response = new ObjectMapper().readValue(rawResponse.getBody(), OAuthResponse.class);


This ObjectMapper is created for each token request. Can we define this as a static final field for this class? Is this readvalue() method thread safe?

On the other hand, calls to this method are not frequent, and soon will be done async.
Specially if readvalue() is not thread safe, I would not bother.

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java

parthban-db · 2025-04-30T21:26:34Z

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java

+    private String audience;
+
+    /**
+     * Validates that a value is non-empty and non-null for required fields.


I think this comment is not correct. We don't have a non-empty check for all the objects. It is only for the Strings. So, we can clarify it here.

parthban-db · 2025-04-30T21:55:38Z

...cks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSourceTest.java

+    if (testCase.expectError) {
+      if (testCase.statusCode == 0) {
+        when(mockHttpClient.execute(any())).thenThrow(new IOException("Network error"));
+      } else {
+        when(mockHttpClient.execute(any()))
+            .thenReturn(
+                new Response(
+                    testCase.responseBody, testCase.statusCode, "Bad Request", new URL(TEST_HOST)));
+      }
+    } else {
+      when(mockHttpClient.execute(any()))
+          .thenReturn(new Response(testCase.responseBody, new URL(TEST_HOST)));
+    }


Can we include the mock HTTP client's request and response in the test case itself?

renaudhartert-db

LGTM overall, deferring final approval to @parthban-db and @hectorcast-db.

hectorcast-db

~~LGTM bar the JavaDocs note from Parth~~

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java

hectorcast-db · 2025-05-02T09:38:18Z

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java

+        OpenIDConnectEndpoints endpoints,
+        IDTokenSource idTokenSource,
+        HttpClient httpClient) {
+      validate(clientId, "ClientID");


To be consistent, we need to validate this on Token(), not when creating the TokenSource.

In other providers, we only check if the fields are set when calling an endpoint/need a token. If we check it earlier, the client construction may fail.

hectorcast-db · 2025-05-06T06:45:01Z

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java

+  @Override
+  public Token getToken() {
+    // Validate all required parameters
+    if (!validate(clientId, "ClientID")


It may be more helpful to the user if they see which parameter is missing. Shall we hvae validate itself throw the exception with the field name in it?
Also note that validate takes the field name, but it does not use it.

parthban-db · 2025-05-06T08:42:26Z

...cks-sdk-java/src/test/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSourceTest.java

+   * required fields cause getToken() to throw IllegalArgumentException.
+   */
+  @Test
+  void testParameterValidation() {


Can this be written as a parameterized test?

parthban-db

LGTM!

github-actions · 2025-05-07T11:46:49Z

If integration tests don't run automatically, an authorized user can run them manually by following the instructions below:

Trigger:
go/deco-tests-run/sdk-java

Inputs:

PR number: 439
Commit SHA: 28dca4f31b963e76cc4a7f3c7cb35573239bb79e

Checks will be approved automatically on success.

Add DatabricksOAuthTokenSource

6f8da14

emmyzhou-db temporarily deployed to test-trigger-is April 30, 2025 15:56 — with GitHub Actions Inactive

emmyzhou-db temporarily deployed to test-trigger-is April 30, 2025 15:57 — with GitHub Actions Inactive

emmyzhou-db requested review from hectorcast-db, parthban-db and renaudhartert-db April 30, 2025 15:59

parthban-db reviewed Apr 30, 2025

View reviewed changes

renaudhartert-db approved these changes May 1, 2025

View reviewed changes

Fix comments and tests

66490d4

emmyzhou-db temporarily deployed to test-trigger-is May 1, 2025 19:22 — with GitHub Actions Inactive

emmyzhou-db temporarily deployed to test-trigger-is May 1, 2025 19:24 — with GitHub Actions Inactive

hectorcast-db approved these changes May 2, 2025

View reviewed changes

emmyzhou-db requested a review from parthban-db May 2, 2025 09:15

parthban-db reviewed May 2, 2025

View reviewed changes

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java Show resolved Hide resolved

databricks-sdk-java/src/main/java/com/databricks/sdk/core/oauth/DatabricksOAuthTokenSource.java Show resolved Hide resolved

hectorcast-db requested changes May 2, 2025

View reviewed changes

Fix comments

de918d3

emmyzhou-db temporarily deployed to test-trigger-is May 2, 2025 10:13 — with GitHub Actions Inactive

emmyzhou-db temporarily deployed to test-trigger-is May 2, 2025 10:14 — with GitHub Actions Inactive

Fix argument validation and add logging

1680580

emmyzhou-db temporarily deployed to test-trigger-is May 5, 2025 17:08 — with GitHub Actions Inactive

Minor comment change

73ffd88

emmyzhou-db temporarily deployed to test-trigger-is May 6, 2025 03:54 — with GitHub Actions Inactive

emmyzhou-db temporarily deployed to test-trigger-is May 6, 2025 03:56 — with GitHub Actions Inactive

emmyzhou-db requested review from hectorcast-db and parthban-db May 6, 2025 03:56

hectorcast-db reviewed May 6, 2025

View reviewed changes

parthban-db reviewed May 6, 2025

View reviewed changes

Fix validation and tests

d843894

emmyzhou-db temporarily deployed to test-trigger-is May 6, 2025 08:54 — with GitHub Actions Inactive

emmyzhou-db temporarily deployed to test-trigger-is May 6, 2025 08:55 — with GitHub Actions Inactive

parthban-db approved these changes May 6, 2025

View reviewed changes

emmyzhou-db requested a review from hectorcast-db May 6, 2025 14:52

emmyzhou-db temporarily deployed to test-trigger-is May 6, 2025 14:53 — with GitHub Actions Inactive

hectorcast-db approved these changes May 6, 2025

View reviewed changes

Merge branch 'main' into emmyzhou-db/db-oauth-token-source

28dca4f

parthban-db temporarily deployed to test-trigger-is May 7, 2025 11:46 — with GitHub Actions Inactive

parthban-db enabled auto-merge May 7, 2025 11:50

parthban-db added this pull request to the merge queue May 7, 2025

Merged via the queue into main with commit efe5aa4 May 7, 2025
15 checks passed

parthban-db deleted the emmyzhou-db/db-oauth-token-source branch May 7, 2025 12:05

Conversation

emmyzhou-db commented Apr 30, 2025

What changes are proposed in this pull request?

How is this tested?

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

parthban-db Apr 30, 2025

Choose a reason for hiding this comment

Uh oh!

hectorcast-db May 2, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

parthban-db Apr 30, 2025

Choose a reason for hiding this comment

Uh oh!

parthban-db Apr 30, 2025

Choose a reason for hiding this comment

Uh oh!

renaudhartert-db left a comment

Choose a reason for hiding this comment

Uh oh!

hectorcast-db left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

hectorcast-db May 2, 2025

Choose a reason for hiding this comment

Uh oh!

hectorcast-db May 6, 2025

Choose a reason for hiding this comment

Uh oh!

parthban-db May 6, 2025

Choose a reason for hiding this comment

Uh oh!

parthban-db left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented May 7, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

hectorcast-db left a comment •

edited

Loading